import pysam
import argparse
import pod5
import numpy as np

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--input', type=str, default="/public/home/hpc234701005/data/R10.4/arab/pod5.CG.features.tsv", required=False,
                        help="input feature tsv")
    parser.add_argument('--output', type=str, default="/public/home/hpc234701005/data/R10.4/arab/train/train.pod5.CG.features.tsv", required=False,
                        help="output feature tsv")
    parser.add_argument('--pod5_dir', type=str, default="/public/data/wangjianxin/nipeng/arab/WHXWZB-2023070072A/raw_data/Nanopore/Col-o/20230723-NPL230985-P6-PAS02934-fast/pod5_pass/train/", required=False,
                        help="id")
    return parser.parse_args()

def read_id(pod5_dir):
    pod5_dr=pod5.DatasetReader(pod5_dir, recursive=True)
    return set(pod5_dr.read_ids)

def deal_tsv(input,output,ids):
    output_fh = open(output, "w", buffering=512)
    with open(input, "r") as f:
        for line in f:
            words = line.strip().split()
            key = words[4]
            if key in ids:
                output_fh.write(line)

if __name__ == '__main__':
    args=parse_args()
    ids=read_id(args.pod5_dir)
    deal_tsv(args.input,args.output,ids)